import os
import numpy as np
import pandas as pd

def calculate_zscore(df):
    if len(df) < 2:
        return np.nan
    mean_score = df['mes_Score'].mean()
    std_score = df['mes_Score'].std()
    if std_score == 0:
        return 0
    return (df['mes_Score'] - mean_score) / std_score

def generate_subject_recommendations():
    chengji = pd.read_csv('data/5_chengji.csv')
    exam_type = pd.read_csv('data/6_exam_type.csv')

    chengji = chengji.merge(exam_type, left_on='exam_type', right_on='EXAM_KIND_ID', how='left')

    chengji = chengji[chengji['mes_Score'] >= 0]

    valid_subjects = ['政治', '历史', '地理', '物理', '化学', '生物', '技术']
    chengji = chengji[chengji['mes_sub_name'].isin(valid_subjects)]

    valid_exam_types = ['十校联考', '五校联考', '期末', '期中']
    chengji = chengji[chengji['EXAM_KIND_NAME'].isin(valid_exam_types)]

    chengji['exam_sdate'] = pd.to_datetime(chengji['exam_sdate'], errors='coerce')
    chengji = chengji.dropna(subset=['exam_sdate'])
    chengji = chengji.sort_values(['mes_StudentID', 'mes_sub_name', 'exam_sdate'])

    exam_stats = chengji.groupby(['mes_TestID', 'mes_sub_id']).apply(
        lambda x: pd.Series({
            'exam_mean': x['mes_Score'].mean(),
            'exam_std': x['mes_Score'].std() if x['mes_Score'].std() != 0 else np.nan
        }),
        include_groups = False
    ).reset_index()

    chengji = chengji.merge(exam_stats, on=['mes_TestID', 'mes_sub_id'], how='left')
    chengji['z_score'] = (chengji['mes_Score'] - chengji['exam_mean']) / chengji['exam_std'].fillna(1)

    student_subject_stats = chengji.groupby(['mes_StudentID', 'mes_sub_name']).agg(
        mean_score=('mes_Score', 'mean'),
        std_score=('mes_Score', 'std'),
        mean_z=('z_score', 'mean'),
        exam_count=('mes_TestID', 'count'),
        last_score=('mes_Score', 'last')
    ).reset_index()

    student_subject_stats = student_subject_stats[student_subject_stats['exam_count'] >= 2]

    for col in ['mean_z', 'mean_score', 'std_score', 'last_score']:
        student_subject_stats[f'{col}_norm'] = (
            student_subject_stats[col] - student_subject_stats[col].min()
        ) / (student_subject_stats[col].max() - student_subject_stats[col].min())

    student_subject_stats['std_score_norm'] = 1 - student_subject_stats['std_score_norm']

    weights = {
        'mean_z_norm': 0.4,
        'mean_score_norm': 0.3,
        'std_score_norm': 0.2,
        'last_score_norm': 0.1
    }

    student_subject_stats['composite'] = (
        student_subject_stats['mean_z_norm'] * weights['mean_z_norm'] +
        student_subject_stats['mean_score_norm'] * weights['mean_score_norm'] +
        student_subject_stats['std_score_norm'] * weights['std_score_norm'] +
        student_subject_stats['last_score_norm'] * weights['last_score_norm']
    )

    recommendations = student_subject_stats.sort_values(
        ['mes_StudentID', 'composite'], 
        ascending=[True, False]
    )

    top3 = recommendations.groupby('mes_StudentID').head(3)
    result = top3.groupby('mes_StudentID')['mes_sub_name'].apply(list).reset_index()
    result[['推荐科目1', '推荐科目2', '推荐科目3']] = result['mes_sub_name'].apply(pd.Series)
    result = result[['mes_StudentID', '推荐科目1', '推荐科目2', '推荐科目3']]
    result.columns = ['学生ID', '推荐科目1', '推荐科目2', '推荐科目3']

    os.makedirs('data', exist_ok=True)

    result.to_csv('data/subject_recommendations.csv', index=False, encoding='utf_8_sig')

if __name__ == '__main__':
    generate_subject_recommendations()